#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/12/18 17:51
# @Author  : CHEN Wang
# @Site    : 
# @File    : asset_ret_decomposition.py
# @Software: PyCharm 

"""
脚本说明： 基金的大类资产收益分解
"""

import pandas as pd
# from quant_researcher.quant.datasource_fetch.fund_api.fund_holding_related import \
#     get_holding_report_info, get_fund_bond_type_allocation
from quant_researcher.quant.datasource_fetch.fund_api.fund_holding_related import \
    get_holding_report_info, get_bond_type_allocation
from quant_researcher.quant.performance_attribution.core_functions.holding_analysis.\
    allocation_analysis.holding_based_allocation import get_asset_allocation
from quant_researcher.quant.datasource_fetch.fund_api.fund_nav_related import get_fund_nav
from quant_researcher.quant.performance_attribution.core_functions.ret_decomposition \
    .asset_ret_decomposition import get_attribution, get_decomposition
from quant_researcher.quant.project_tool.QACode import stock_code_cleaning
from quant_researcher.quant.performance_attribution.core_functions.attribution_analysis import multi_period_attribution


def fund_asset_decomposition(fund_code='110022', benchmark='TK30B70S',
                             start_date='2019-12-31', end_date='2020-06-30', **kwargs):
    """
    单期基金大类资产收益率分解（权益，债券，现金，其他），需要基金的全部持仓信息，因此每半年计算一次.
    注：大类资产权重应该是净资产占比

    这期间
    股票持仓认为是两端股票持仓的平均权重，并乘以每个股票这段期间内收益
    债券持仓认为是两端债券券种配置的平均权重，并乘以每类债券券种对应代理指数的区间内收益
    现金持仓认为是两端现金及现金等价物部分的平均权重，并乘以货币基金指数的区间收益
    其他部分的收益为基金区间收益 - 上述三部分的收益率之和

    :param str fund_code: 基金代码，如："110022"
    :param str benchmark: 基准代码，如："TK30B70S"，目前仅支持混合基准
    :param str start_date: 开始日期，如："2020-06-30"，大类资产归因分析基于半年报和年报数据，
        开始日期和结束日期必须为半年报，年报截止日期，如："2020-06-30"，”2019-12-31“
    :param str end_date: 结束日期，如："2020-06-30"，大类资产归因分析基于半年报和年报数据，
        开始日期和结束日期必须为半年报，年报截止日期，如："2020-06-30"，”2019-12-31“
    :param kwargs:
        - conn_base: 基础数据库连接
        - conn_ty: 投研数据库连接
        - conn_stock: 股票数据库连接
    :return: 收益归因和收益分解

    """
    conn_base = kwargs.pop('conn_base', None)
    conn_ty = kwargs.pop('conn_ty', None)
    conn_stock = kwargs.pop('conn_stock', None)
    # 获取现金的期初期末配置权重
    asset_weight_start = get_asset_allocation([fund_code], net_or_total='n', end_date=start_date, conn=conn_base)
    cash_weight_start = asset_weight_start['cash'].iloc[0]
    asset_weight_end = get_asset_allocation([fund_code], net_or_total='n', end_date=end_date, conn=conn_base)
    cash_weight_end = asset_weight_end['cash'].iloc[0]
    cash_df_start = pd.DataFrame({'asset_type': ['cash'], 'security_code': ['H11025'],
                                  'security_weight': [cash_weight_start]})
    cash_df_end = pd.DataFrame({'asset_type': ['cash'], 'security_code': ['H11025'],
                                'security_weight': [cash_weight_end]})

    # 获取股票的期初期末配置权重
    stock_df_start = get_holding_report_info(report='stock', fund_code=[fund_code], report_type=[5, 6],
                                             end_date=start_date,
                                             select=['fund_code', 'stock_code',
                                                     'hold_mktvalue as value',
                                                     'pptinnv as stock_weight'], conn=conn_base)
    if stock_df_start is None:   # 如果没有股票持仓
        stock_df_start = pd.DataFrame({'security_code': ['000001'], 'security_weight': [0]})
    else:
        stock_df_start['stock_code'] = stock_code_cleaning(stock_df_start['stock_code'])
        stock_df_start = stock_df_start.dropna()
        stock_df_start['stock_weight'] = stock_df_start['stock_weight'] / 100
        stock_df_start = stock_df_start[['stock_code', 'stock_weight']].rename(
            columns={'stock_code': 'security_code', 'stock_weight': 'security_weight'})
    stock_df_start['asset_type'] = 'stock'

    stock_df_end = get_holding_report_info(report='stock', fund_code=[fund_code], report_type=[5, 6],
                                           end_date=end_date,
                                           select=['fund_code', 'stock_code',
                                                   'hold_mktvalue as value',
                                                   'pptinnv as stock_weight'], conn=conn_base)
    if stock_df_end is None:  # 如果没有股票持仓
        stock_df_end = pd.DataFrame({'security_code': ['000001'], 'security_weight': [0]})
    else:
        stock_df_end['stock_code'] = stock_code_cleaning(stock_df_end['stock_code'])
        stock_df_end = stock_df_end.dropna()
        stock_df_end['stock_weight'] = stock_df_end['stock_weight'] / 100
        stock_df_end = stock_df_end[['stock_code', 'stock_weight']].rename(
            columns={'stock_code': 'security_code', 'stock_weight': 'security_weight'})
    stock_df_end['asset_type'] = 'stock'

    # 获取券种的期初期末配置权重, 如果没有债券持仓, get_fund_bond_type_allocation会返回一个权重为0的数据
    bond_index_dict = {'rate_bonds': 'CBA00601.CS', 'credit_bonds': 'CBA02701.CS',
                       'conver_bonds': '000832'}
    # bond_df_start = get_fund_bond_type_allocation([fund_code], 'quarterly', end_date=start_date)
    bond_df_start = get_bond_type_allocation(fund_code, only_latest=True, end_date=start_date)
    if bond_df_start is None:
        bond_df_start = pd.DataFrame({'bond_type': ['rate_bonds', 'credit_bonds', 'conver_bonds']})
        bond_df_start['rept_enddate'] = end_date
        bond_df_start['weight'] = 0
    else:
        bond_df_start = bond_df_start.rename(columns={'report_date': 'rept_enddate'})
        bond_df_start = bond_df_start.drop(columns=['fund_code'])
        bond_df_start['rept_enddate'] = bond_df_start['rept_enddate'].apply(lambda x: x.replace('-', ''))
    bond_df_start['security_code'] = bond_df_start['bond_type'].replace(bond_index_dict)
    bond_df_start = bond_df_start[['security_code', 'weight']].rename(
        columns={'weight': 'security_weight'})
    bond_df_start['asset_type'] = 'bond'
    # bond_df_end = get_fund_bond_type_allocation([fund_code], 'quarterly', end_date=end_date)
    bond_df_end = get_bond_type_allocation(fund_code, only_latest=True, end_date=end_date)
    if bond_df_end is None:
        bond_df_end = pd.DataFrame({'bond_type': ['rate_bonds', 'credit_bonds', 'conver_bonds']})
        bond_df_end['rept_enddate'] = end_date
        bond_df_end['weight'] = 0
    else:
        bond_df_end = bond_df_end.rename(columns={'report_date': 'rept_enddate'})
        bond_df_end = bond_df_end.drop(columns=['fund_code'])
        bond_df_end['rept_enddate'] = bond_df_end['rept_enddate'].apply(lambda x: x.replace('-', ''))
    bond_df_end['security_code'] = bond_df_end['bond_type'].replace(bond_index_dict)
    bond_df_end = bond_df_end[['security_code', 'weight']].rename(
        columns={'weight': 'security_weight'})
    bond_df_end['asset_type'] = 'bond'

    # 获取其他资产的期初期末配置权重
    other_weight_start = asset_weight_start['other'].iloc[0]
    other_weight_end = asset_weight_end['other'].iloc[0]
    other_df_start = pd.DataFrame({'asset_type': ['other'], 'security_code': ['other'],
                                   'security_weight': [other_weight_start]})
    other_df_end = pd.DataFrame({'asset_type': ['other'], 'security_code': ['other'],
                                 'security_weight': [other_weight_end]})

    # 期初期末资产配置权重合并
    asset_df_start = pd.concat([cash_df_start, stock_df_start, bond_df_start, other_df_start],
                               ignore_index=True, sort=False)
    asset_df_end = pd.concat([cash_df_end, stock_df_end, bond_df_end, other_df_end],
                             ignore_index=True, sort=False)
    asset_df_start['date'] = start_date
    asset_df_end['date'] = end_date

    # 获取基金区间收益率
    nav_df = get_fund_nav(fund_code=fund_code, start_date=start_date, end_date=end_date, conn=conn_ty)
    nav_df['end_date'] = pd.to_datetime(nav_df['end_date'])
    nav_df['end_date'] = nav_df['end_date'].dt.date
    total_ret = nav_df['adjusted_nav'].iloc[-1] / nav_df['adjusted_nav'].iloc[0] - 1

    # 对基金区间收益率进行分解
    decomp = get_decomposition(holding_begin=asset_df_start,
                               holding_end=asset_df_end,
                               total_return=total_ret,
                               conn_base=conn_base, conn_stock=conn_stock)

    # 对基金区间主动收益率进行大类brinson归因
    attribu = get_attribution(asset_df=decomp, begin_date=start_date, end_date=end_date,
                              benchmark=benchmark, conn=conn_base)

    result = decomp.merge(attribu, on='asset_type', how='left')
    result['end_date'] = end_date

    return result, decomp, attribu


def fund_ret_decomposition_muti_period(fund_code='110022', start_date='2019-12-31',
                                       end_date='2020-06-30', benchmark='TK30B70S', **kwargs):
    conn_base = kwargs.pop('conn_base', None)
    conn_ty = kwargs.pop('conn_ty', None)
    conn_stock = kwargs.pop('conn_stock', None)

    date_range = pd.date_range(start_date, end_date, freq='Q')
    date_to_calc = date_range[date_range.month.isin([6, 12])].strftime('%Y-%m-%d').tolist()
    result_list = []
    for start, end in zip(date_to_calc[:-1], date_to_calc[1:]):
        result, decomp, attribu = fund_asset_decomposition(fund_code, benchmark, start, end,
                                                           conn_base=conn_base, conn_stock=conn_stock, conn_ty=conn_ty)
        result_list.append(result)
    full_result = pd.concat(result_list)
    asset_type_dict = {'stock': '10', 'bond': '20', 'cash': '30', 'other': '40'}
    full_result['asset_type'] = full_result['asset_type'].replace(asset_type_dict)

    # 多期大类收益率归因
    col_needed_decomp = ['end_date', 'all_allocation_ability', 'all_manage_ability',
                         'all_other_ability', 'total_return_p', 'total_return_b']
    attribution_df = full_result[col_needed_decomp].drop_duplicates(subset=['end_date'])
    attribution_df = attribution_df.set_index('end_date')

    period_attribution_list = []
    for date in attribution_df.index:
        attribution_df_temp = attribution_df.loc[:date, ['all_allocation_ability', 'all_manage_ability',
                         'all_other_ability']]
        fund_ret_series = attribution_df.loc[:date, 'total_return_p']
        benchmark_ret_series = attribution_df.loc[:date, 'total_return_b']
        fund_cum_return = ((1 + fund_ret_series).cumprod() - 1)[-1]
        benchmark_cum_return = ((1 + benchmark_ret_series).cumprod() - 1)[-1]
        cum_active_return = fund_cum_return - benchmark_cum_return
        multi_period_coefficient = multi_period_attribution.relative_multi_period(
            fund_ret_series, benchmark_ret_series)  # 获取多期归因系数

        weighted_ret = attribution_df_temp.mul(multi_period_coefficient, axis=0)
        cummulated_ret = weighted_ret.sum(axis=0)
        cummulated_ret['end_date'] = date
        cummulated_ret['fund_cum_return'] = fund_cum_return
        cummulated_ret['benchmark_cum_return'] = benchmark_cum_return
        cummulated_ret['cum_active_ret'] = cum_active_return

        period_attribution_list.append(cummulated_ret)
    period_attribution = pd.DataFrame(period_attribution_list).set_index(['end_date'])
    period_attribution.loc[start_date, :] = 0
    period_attribution.sort_index(inplace=True)

    # 获取从起始日至今的累计收益率归因
    latest_attribution = period_attribution.iloc[-1]

    # 多期大类收益率分解
    ret_decomp = full_result.set_index(['end_date', 'asset_type'])['weighted_return_x'].unstack()
    fund_return_each_period = full_result.groupby('end_date')['total_return_p'].first()
    ret_decomp['fund_return_each_period'] = fund_return_each_period
    period_decomposition_list = []
    for date in ret_decomp.index:
        ret_decomp_temp = ret_decomp.loc[:date, ['10', '20', '30', '40']]
        fund_ret_series = ret_decomp.loc[:date, 'fund_return_each_period']
        fund_cum_return = ((1 + fund_ret_series).cumprod() - 1)[-1]
        multi_period_coefficient = multi_period_attribution.absolute_multi_period(
                                fund_ret_series)  # 获取多期归因系数
        weighted_ret = ret_decomp_temp.mul(multi_period_coefficient, axis=0)
        cummulated_ret = weighted_ret.sum(axis=0)
        cummulated_ret['end_date'] = date
        cummulated_ret['fund_cum_return'] = fund_cum_return
        period_decomposition_list.append(cummulated_ret)

    period_decomposition = pd.DataFrame(period_decomposition_list).set_index(['end_date'])
    period_decomposition.loc[start_date, :] = 0
    period_decomposition.sort_index(inplace=True)

    # 获取从起始日至今的累计收益率分解
    latest_decomp = period_decomposition.iloc[-1]

    return latest_attribution, latest_decomp, period_decomposition


if __name__ == '__main__':
    fund_code = '110022'  # 易方达消费行业 -- 用来测试即有股又有债的基金
    # fund_code = '002354'  # 博时裕腾纯债   -- 用来测试只有债券的基金
    benchmark = 'TK30B70S'  # 业绩基准- 30%债券70%股票

    # 测试 fund_asset_decomposition
    # start_date = '2019-12-31'
    # end_date = '2020-06-30'
    # test = fund_asset_decomposition(fund_code=fund_code, benchmark=benchmark,
    #                                 start_date=start_date, end_date=end_date)

    # 测试 fund_ret_decomposition_muti_period
    start_date = '2020-12-31'
    end_date = '2021-06-30'
    test = fund_ret_decomposition_muti_period(fund_code=fund_code, benchmark=benchmark,
                                              start_date=start_date, end_date=end_date)
